package com.taoz.rpc.core.registry;

import cn.hutool.json.JSONUtil;
import com.taoz.rpc.core.config.ApplicationConfig;
import com.taoz.rpc.core.model.ServiceMetaInfo;
import io.etcd.jetcd.*;
import io.etcd.jetcd.options.GetOption;
import io.etcd.jetcd.options.PutOption;
import io.etcd.jetcd.options.WatchOption;
import io.vertx.core.impl.ConcurrentHashSet;
import lombok.extern.slf4j.Slf4j;

import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Set;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

@Slf4j
public class EtcdRegistry implements Registry {
    private Client client;
    private KV kvClient;
    private Lease leaseClient;
    private final Set<ServiceMetaInfo> serviceSet = new CopyOnWriteArraySet<>();
    private final RegistryServiceCache serviceCache = new RegistryServiceCache();
    private final Set<String> watchedService = new ConcurrentHashSet<>();

    @Override
    public void init() {
        ApplicationConfig config = ApplicationConfig.getConfig();
        client = Client.builder().endpoints(config.getRegistryAddress()).build();
        kvClient = client.getKVClient();
        leaseClient = client.getLeaseClient();
        if (config.isServer()) heatBeat();
        Runtime.getRuntime().addShutdownHook(new Thread(this::destroy));
    }

    @Override
    public void register(ServiceMetaInfo info) throws Exception {
        ByteSequence key = ByteSequence.from(info.getServiceKey(), StandardCharsets.UTF_8);
        ByteSequence value = ByteSequence.from(JSONUtil.toJsonStr(info), StandardCharsets.UTF_8);;
        long leaseId = leaseClient.grant(60).get().getID();
        PutOption putOption = PutOption.builder().withLeaseId(leaseId).build();
        kvClient.put(key, value, putOption).get();
        serviceSet.add(info);
    }

    @Override
    public void deregister(ServiceMetaInfo info) throws Exception {
        ByteSequence key = ByteSequence.from(info.getServiceKey(), StandardCharsets.UTF_8);
        kvClient.delete(key).get();
        serviceSet.remove(info);
    }

    @Override
    public List<ServiceMetaInfo> serviceDiscovery(ServiceMetaInfo info) throws Exception {
        String name = info.getName();
        List<ServiceMetaInfo> serviceList = serviceCache.getService(name);
        if (serviceList != null) {
            return serviceList;
        }

        log.info("正在从配置中心拉取服务" + name);
        String servicePrefix = String.format("/%s/%s/", ApplicationConfig.getConfig().getName(), info.getName());
        GetOption getOption = GetOption.builder().isPrefix(true).build();
        List<KeyValue> kvs = kvClient.get(ByteSequence.from(servicePrefix, StandardCharsets.UTF_8), getOption).get().getKvs();
        serviceList = kvs.stream().map(e ->
                JSONUtil.toBean(e.getValue().toString(StandardCharsets.UTF_8), ServiceMetaInfo.class)
        ).toList();

        watch(servicePrefix, name);

        serviceCache.cacheService(name, serviceList);
        return serviceList;
    }

    @Override
    public void heatBeat() {
        ScheduledExecutorService executor = Executors.newSingleThreadScheduledExecutor();
        executor.scheduleAtFixedRate(() -> {
            log.info("正在进行心跳续期");
            for (ServiceMetaInfo info : serviceSet) {
                try {
                    register(info);
                } catch (Exception e) {
                    log.info("续期失败", e);
                }
            }
        }, 20, 20, TimeUnit.SECONDS);
    }

    @Override
    public void watch(String servicePrefix, String serviceName) {
        if (watchedService.contains(servicePrefix)) {
            return;
        }
        watchedService.add(servicePrefix);

        log.info("正在监听" + servicePrefix);
        client.getWatchClient().watch(ByteSequence.from(servicePrefix, StandardCharsets.UTF_8),
                WatchOption.builder().isPrefix(true).build(),
                resp -> {
                    log.info("监测到服务 {} 发生变更，正在移除旧服务", serviceName);
                    serviceCache.removeService(serviceName);}
        );
    }

    @Override
    public void destroy() {
        log.info("正在关闭 Etcd Registry服务");
        for (ServiceMetaInfo info : serviceSet) {
            try {
                deregister(info);
            } catch (Exception e) {
                log.info("销毁服务 {} 失败", info);
            }
        }
        if (client != null) client.close();
        if (kvClient != null) kvClient.close();
        if (leaseClient != null) leaseClient.close();
    }
}
